package com.surprise.db.datasources;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;
import java.util.Random;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.jdbc.datasource.lookup.DataSourceLookup;
import org.springframework.util.Assert;

/**
 *
 * 2013-11-17
 */
public class AutoDataSource extends AbstractRoutingDataSource {
    private Object[] targetDataSourcesKeys = null;

    private static final ThreadLocal<Object> contextHolder = new ThreadLocal<Object>();

    public static void setCurrentLookupKeyIndex(Integer currentLookupKeyIndex) {
        Assert.notNull(currentLookupKeyIndex,
                "current LookupKey cannot be null");
        contextHolder.set(currentLookupKeyIndex);
    }

    public static Integer getCurrentLookupKeyIndex() {
        Integer i = (Integer) contextHolder.get();
        if (i == null) {
            setCurrentLookupKeyIndex(0);
            return Integer.valueOf(0);
        } else {
            return i;
        }
    }

    public static void clearCurrentLookupKeyIndex() {
        contextHolder.remove();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return targetDataSourcesKeys[getCurrentLookupKeyIndex().intValue()];
    }

    @Override
    public Connection getConnection() throws SQLException {
        int startIndex = new Random().nextInt(targetDataSourcesKeys.length);
        setCurrentLookupKeyIndex(startIndex % targetDataSourcesKeys.length);
        for (int i = 0; i < targetDataSourcesKeys.length;) {
            try {
                Connection conn = determineTargetDataSource().getConnection();
                // 获取连接成功
                return conn;
            } catch (Exception e) {
                // 获取连接失败
                logger.warn("", e);
                logger.warn("Can't get the connection from:"
                        + determineTargetDataSource().getConnection()
                                .toString());
                i++;
                setCurrentLookupKeyIndex((startIndex + i)
                        % targetDataSourcesKeys.length);
            }
        }
        try {
            return determineTargetDataSource().getConnection();
        } catch (SQLException e) {
            logger.error("", e);
            throw e;
        }
    }

    @Override
    public void setDataSourceLookup(DataSourceLookup dataSourceLookup) {
        super.setDataSourceLookup(dataSourceLookup);
    }

    @Override
    public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
        super.setDefaultTargetDataSource(defaultTargetDataSource);
    }

    @Override
    public void setTargetDataSources(Map targetDataSources) {
        this.targetDataSourcesKeys = targetDataSources.keySet().toArray();
        setCurrentLookupKeyIndex(0);
        super.setTargetDataSources(targetDataSources);
    }

}
